import psycopg2
from psycopg2 import sql
import os
from dotenv import load_dotenv
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 加载环境变量
load_dotenv()

class PostgreSQLDatabase:
    def __init__(self):
        """
        初始化数据库连接配置
        """
        self.host = os.getenv('DB_HOST', 'localhost')
        self.database = os.getenv('DB_NAME', 'testdb')
        self.user = os.getenv('DB_USER', 'postgres')
        self.password = os.getenv('DB_PASSWORD', 'password')
        self.port = os.getenv('DB_PORT', '5432')
        self.connection = None
    
    def connect(self):
        """
        建立数据库连接
        """
        try:
            self.connection = psycopg2.connect(
                host=self.host,
                database=self.database,
                user=self.user,
                password=self.password,
                port=self.port
            )
            logger.info("成功连接到PostgreSQL数据库!")
            return True
        except Exception as error:
            logger.error(f"连接PostgreSQL时出错: {error}")
            return False
    
    def disconnect(self):
        """
        断开数据库连接
        """
        if self.connection:
            self.connection.close()
            logger.info("数据库连接已关闭!")
    
    def execute_query(self, query, params=None):
        """
        执行查询语句
        """
        try:
            cursor = self.connection.cursor()
            if params:
                cursor.execute(query, params)
            else:
                cursor.execute(query)
            
            if query.strip().upper().startswith('SELECT'):
                result = cursor.fetchall()
                cursor.close()
                return result
            else:
                self.connection.commit()
                cursor.close()
                return True
        except Exception as error:
            logger.error(f"执行查询时出错: {error}")
            self.connection.rollback()
            return None
    
    def execute_many(self, query, params_list):
        """
        执行批量操作
        """
        try:
            cursor = self.connection.cursor()
            cursor.executemany(query, params_list)
            self.connection.commit()
            cursor.close()
            return True
        except Exception as error:
            logger.error(f"执行批量操作时出错: {error}")
            self.connection.rollback()
            return False
    
    def create_table(self, table_name, columns):
        """
        创建表
        :param table_name: 表名
        :param columns: 列定义字典，格式: {'列名': '数据类型 约束'}
        """
        try:
            column_defs = []
            for col_name, col_def in columns.items():
                column_defs.append(f"{col_name} {col_def}")
            
            create_query = f"""
            CREATE TABLE IF NOT EXISTS {table_name} (
                {', '.join(column_defs)}
            );
            """
            
            cursor = self.connection.cursor()
            cursor.execute(create_query)
            self.connection.commit()
            cursor.close()
            logger.info(f"表 '{table_name}' 创建成功!")
            return True
        except Exception as error:
            logger.error(f"创建表时出错: {error}")
            return False